from typing import Dict, List, Set, Tuple
import itertools
from collections import namedtuple
from sklearn.metrics import precision_score, recall_score, f1_score
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoConfig,
    AutoModel,
    AutoTokenizer,
    BartTokenizer,
    BertModel,
    BertTokenizer,
    AutoModelForMaskedLM,
)

CHINESE_PUNCTUATIONS = "！？。＂＃＄％＆＇（）＊+＋，-－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
ENGLISH_PUNCTUATIONS = "!\"#$%&'()*+?,-./:;<=>?@[\]^_`{|}~"
PUNCTUATIONS = CHINESE_PUNCTUATIONS + ENGLISH_PUNCTUATIONS

CONFIG_MODEL_TOKENIZER_CLASSES = {
    "auto-model": (AutoConfig, AutoModel, AutoTokenizer),
    "fnlp/bart-base-chinese": (AutoConfig, AutoModelForSeq2SeqLM, BertTokenizer),
    "uer/bart-base-chinese-cluecorpussmall": (AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer),
    "hfl/chinese-roberta-wwm-ext": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "hfl/chinese-bert-wwm-ext": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "hfl/chinese-macbert-base": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "hfl/chinese-macbert-large": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "hfl/chinese-roberta-wwm-ext-large": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "freedomking/mc-bert": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "trueto/medbert-base-wwm-chinese": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "trueto/medbert-kd-chinese": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
    "yechen/bert-large-chinese": (AutoConfig, AutoModelForMaskedLM, AutoTokenizer),
}

def simple_metric_score(y_true: List[int], y_pred: List[int]) -> Dict:

    p = precision_score(y_true, y_pred, zero_division=0)
    r = recall_score(y_true, y_pred, zero_division=0)
    f1 = f1_score(y_true, y_pred, zero_division=0)
    return {'f1_score': f1,
            'precision':p,
            'recall': r,
            }
# 实体集
Entity = namedtuple('Entity', ['type', 'entity'])
# 这里是使用的Macro计算方法
def metric_score(true_entities: Set[Tuple], pred_entities: Set[Tuple]):
    """Compute the F1 score."""
    nb_correct = len(true_entities & pred_entities) # TP
    nb_pred = len(pred_entities) # TP + FP
    nb_true = len(true_entities) # TP + FN

    p = nb_correct / nb_pred if nb_pred > 0 else 0
    r = nb_correct / nb_true if nb_true > 0 else 0
    f1 = 2 * p * r / (p + r) if p + r > 0 else 0

    return {'f1_score': f1,
            'precision':p,
            'recall': r,
            'num_correct': nb_correct,
            'num_pred': nb_pred,
            'num_true': nb_true
            }

class _Class:
    
    
    def __init__(self):
        self.MASK_TOKEN = '[MASK]'
        self.sep_token = '[SEP]'
        self.cls_token = '[CLS]'

        self.num_lbl = 3
        
        self.pet_patterns = [
            
        ]
    
    @staticmethod
    def func1(_entity: str, _type: str) -> str:
        return ''
    
    @classmethod
    def func2(cls, idx):
        return cls().pet_patterns[idx]
        
    @property
    def func3(self):
        return self.num_lbl

    


